from pymilvus import utility, FieldSchema, DataType, CollectionSchema, Collection

# 定义Milvus中用于存储文档向量的集合名称
DOC_COLLECTION_NAME = "doc_vectors"
def init_doc_collection():
    # 如果Milvus中不存在指定名称的集合，则进行创建
    if not utility.has_collection(DOC_COLLECTION_NAME):
        #定义集合的字段，包括自增主键id,浮点向量embedding和原文档内容doc
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=2560),
            FieldSchema(name="doc", dtype=DataType.VARCHAR, max_length=4096),
        ]
        # 根据字段定义创建集合的schema
        schema = CollectionSchema(fields=fields, description="文档向量集合")
        # 创建集合对象
        col = Collection(DOC_COLLECTION_NAME, schema=schema)

        # 定义向量字段的索引参数
        index_params = {
            "index_type": "IVF_FLAT",
            "metric_type": "COSINE",
            "params": {"nlist": 128}
        }

        # 为embedding字段创建索引
        col.create_index(field_name="embedding", index_params=index_params)
        # 加载集合到内存
        col.load()
    else:
        col = Collection(DOC_COLLECTION_NAME)
        col.load()

def insert_doc_vectors(embeddings, docs):
    col = Collection(DOC_COLLECTION_NAME)
    # 组织插入的数据，顺序与schema一致
    data = [embeddings, docs]
    # 插入数据
    col.insert(data)

def insert_doc_vector(embedding, doc):
    col = Collection(DOC_COLLECTION_NAME)
    # 组织插入的数据，顺序与schema一致
    data = {
        "embedding": embedding,
        "doc": doc
    }
    # 插入数据
    col.insert(data)

def search_doc_vectors(query_embedding, top_k):
    col = Collection(DOC_COLLECTION_NAME)
    results = col.search(
        data=[query_embedding],
        anns_field="embedding",
        param={
            "metric_type": "COSINE",
            "params": {
                "nprobe": 10
            }
        },
        limit=top_k,
        output_fields=["doc"]
    )
    return results